from openai import OpenAI
import os
from dotenv import load_dotenv

# pip install openai dotenv


class DeepseekService:
    def __init__(self):
        """初始化Deepseek服务实例"""
        self.api_key = self.get_api_key()
        self.base_url = self.get_base_url()
        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)

    def get_api_key(self):
        """获取API KEY"""
        load_dotenv()
        api_key = os.getenv("DEEPSEEK_API_KEY")
        return api_key

    def get_base_url(self, platform="deepseek"):
        """获取BASE URL"""
        if platform == "deepseek":
            # print("已经获取到了BASE_URL，平台是deepseek")
            return "https://api.deepseek.com"
        else:
            print("不支持的平台")
            return None

    def ai_chat(self, prompt, **kwargs):
        """
        调用AI接口进行对话
        :param prompt: 提问内容
        :param kwargs: 其他参数
        :return: AI回答内容
        """

        prompts = f"""
        你是一个人工智能助手，你的任务是回答用户的问题。
        请用中文回答，不要用英文回答。
        用户的问题如下：{prompt}
        """

        stream = kwargs.get("stream", False)
        response = self.client.chat.completions.create(
            model=kwargs.get("model", "deepseek-chat"),
            messages=[
                {"role": "system", "content": "You are a helpful assistant"},
                {"role": "user", "content": f"{prompts}"},
            ],
            temperature=kwargs.get("temperature", 1.0),
            max_tokens=kwargs.get("max_tokens", 4096),
            stream=stream,
        )

        if stream:
            contents = ""
            print("\nAI回复：")
            for chunk in response:
                content = chunk.choices[0].delta.content
                print(content, end="", flush=True)
                contents += content
        else:
            contents = response.choices[0].message.content
            print("\nAI回复：")
            print(contents)

        return contents


if __name__ == "__main__":
    # 自定义模型参数
    custom_params = {"temperature": 0.9, "max_tokens": 150, "stream": True}

    deepseekservice = DeepseekService()

    while True:
        prompt = input("\n请输入您的问题：")

        if prompt.lower() == "exit":
            print("程序已退出")
            break

        # 通过**解包字典作为关键字参数
        response_content = deepseekservice.ai_chat(prompt, **custom_params)
        print(f"AI回复：{response_content}")
